import numpy as np
import deepxde as dde
import matplotlib.pyplot as plt
from abc import abstractmethod
from deepxde.backend import backend_name
from deepxde.backend import torch, paddle, tf, jax
from abc import ABC, abstractmethod

class cFPCT_DNN_Cases(ABC):
    def __init__(self,
                name, v, 
                NumTrain, 
                NumTest, 
                max_iter, 
                FP_size, 
                DNN_size):
        self.v = v
        self.NumTrain = NumTrain
        self.NumTest = NumTest
        self.max_iter = max_iter
        self.FP_size = FP_size
        self.DNN_size = DNN_size
        self.name = name
        self.backend = self.gen_backend()
        self.pde = self.gen_pde()
        self.geomtime = self.gen_geomtime()
        self.data = self.gen_data()

    def gen_backend(self):
        if backend_name == 'pytorch':
            return torch
        elif backend_name == 'paddle':
            return paddle
        elif backend_name in ['tensorflow', 'tensorflow.compat.v1']:
            print('Currently tensorflow backend has not been tested, if you find any bug, please contact us.')
            return tf.math
        elif backend_name == 'jax':
            print('Currently jax backend has not been tested, if you find any bug, please contact us.')
            return jax.numpy
        else:
            raise Warning('Currently only support pytorch and paddle backend')

    @abstractmethod
    def func(self, x):
        pass

    def gen_pde(self):
        pass
    
    @abstractmethod
    def gen_geomtime(self):
        pass

    def gen_data(self):
        return dde.data.Function(self.geomtime, self.func, num_train=self.NumTrain, num_test=self.NumTest, train_distribution='uniform')
    
    def gen_net(self, layer_size, activation, initializer):
        net = dde.maps.FNN(layer_size, activation, initializer)
        if self.use_output_transform:
            net.apply_output_transform(self.output_transform)
        return net
    
    def gen_testdata(self):
        if callable(self.func):
            x = self.geomtime.random_points(self.NumTest)
            y = self.func(x)
            return x, y
        else:
            raise Warning('You must rewrite one of func() and gen_testdata()')
    
    def plot(self, model=None, title='Solution'):
        pass


class HighFrequencyCase1(cFPCT_DNN_Cases):
    def __init__(self,
                 v=3.0, 
                 NumTrain=1000, 
                 NumTest=2000, 
                 max_iter=40000, 
                 FP_size=[1] + [100] + [1], 
                 DNN_size=[1] + [10] * 3 + [1]):
        super().__init__(name="HighFrequencyCase1", v=v, NumTrain=NumTrain, NumTest=NumTest, max_iter=max_iter, FP_size=FP_size, DNN_size= DNN_size)

    def func(self, x):
        return np.sin(50*x) + np.cos(40*x)

    def gen_geomtime(self):
        return dde.geometry.Interval(-np.pi, np.pi)
    
    def plot(self, model=None, title='Solution'):
        x = np.linspace(-np.pi, np.pi, 2000)
        x = np.expand_dims(x, axis=1)
        y = self.func(x)
        plt.plot(x, y, label="Ground truth", alpha=0.5)

        if model is not None:
            y_pred = model.model.predict(x)
            plt.plot(x, y_pred, label="Prediction", alpha=0.5)
            plt.plot(x, y_pred - y, label="Error", alpha=0.5)
        
        plt.title(title)
        plt.xlabel("x")
        plt.ylabel("y")
        plt.legend()
        plt.show()

class HighFrequencyCase2(cFPCT_DNN_Cases):
    def __init__(self,
                 v=6.0, 
                 NumTrain=320, 
                 NumTest=2000, 
                 max_iter=40000, 
                 FP_size=[1] + [200] + [1], 
                 DNN_size=[1] + [20] * 3 + [1]):
        super().__init__(name="HighFrequencyCase2", v=v, NumTrain=NumTrain, NumTest=NumTest, max_iter=max_iter, FP_size=FP_size, DNN_size= DNN_size)

    def func(self, x):
        return np.sin(60*x)

    def gen_geomtime(self):
        return  dde.geometry.Interval(-np.pi, np.pi)
    
    def plot(self, model=None, title='Solution'):
        x = np.linspace(-np.pi, np.pi, 2000)
        x = np.expand_dims(x, axis=1)
        y = self.func(x)
        plt.plot(x, y, label="Ground truth", alpha=0.5)

        if model is not None:
            y_pred = model.model.predict(x)
            plt.plot(x, y_pred, label="Prediction", alpha=0.5)
            plt.plot(x, y_pred - y, label="Error", alpha=0.5)
        
        plt.title(title)
        plt.xlabel("x")
        plt.ylabel("y")
        plt.legend()
        plt.show()

class HighFrequencyCase3(cFPCT_DNN_Cases):
    def __init__(self,
                 v=6.0, 
                 NumTrain=320, 
                 NumTest=2000, 
                 max_iter=40000, 
                 FP_size=[1] + [200] + [1], 
                 DNN_size=[1] + [20] * 3 + [1]):
        super().__init__(name="HighFrequencyCase3", v=v, NumTrain=NumTrain, NumTest=NumTest, max_iter=max_iter, FP_size=FP_size, DNN_size= DNN_size)

    def func(self, x):
        return np.sin(100*x) + np.sin(160*x)

    def gen_geomtime(self):
        return dde.geometry.Interval(0.5, np.pi)
    
    def plot(self, model=None, title='Solution'):
        x = np.linspace(0.5, np.pi, 2000)
        x = np.expand_dims(x, axis=1)
        y = self.func(x)
        plt.plot(x, y, label="Ground truth", alpha=0.5)

        if model is not None:
            y_pred = model.model.predict(x)
            plt.plot(x, y_pred, label="Prediction", alpha=0.5)
            plt.plot(x, y_pred - y, label="Error", alpha=0.5)
        
        plt.title(title)
        plt.xlabel("x")
        plt.ylabel("y")
        plt.legend()
        plt.show()

class AKG(cFPCT_DNN_Cases):
    def __init__(self,
                 improved=False,
                 k1=7.0,
                 k2=12.0,
                 v=5.0, 
                 NumTrain=300, 
                 NumTest=2000, 
                 max_iter=60000, 
                 FP_size=[1] + [200] + [1], 
                 DNN_size=[1] + [20] * 3 + [1]):
        if improved:
            k1=27.0
            k2=50.0
            v=6.0
            FP_size=[1] + [260] + [1] 
        self.k1 = k1
        self.k2 = k2
        super().__init__(name="AKG", v=v, NumTrain=NumTrain, NumTest=NumTest, max_iter=max_iter, FP_size=FP_size, DNN_size= DNN_size)

    def func(self, x):
        return (x ** 3 - x) * np.sin(self.k1 * x) / self.k1 + np.sin(self.k2 * x)

    def gen_geomtime(self):
        return  dde.geometry.Interval(-3.0, 3.0)
    
    def plot(self, model=None, title='Solution'):
        x = np.linspace(-3.0, 3.0, 2000)
        x = np.expand_dims(x, axis=1)
        y = self.func(x)
        plt.plot(x, y, label="Ground truth", alpha=0.5)

        if model is not None:
            y_pred = model.model.predict(x)
            plt.plot(x, y_pred, label="Prediction", alpha=0.5)
            plt.plot(x, y_pred - y, label="Error", alpha=0.5)
        
        plt.title(title)
        plt.xlabel("x")
        plt.ylabel("y")
        plt.legend()
        plt.show()

class Helmholtz(cFPCT_DNN_Cases):
    def __init__(self,
                 improved=False,
                 a=0.5,
                 b=3.14,
                 k=20.0,
                 v=5.0, 
                 NumTrain=300, 
                 NumTest=2000, 
                 max_iter=60000, 
                 FP_size=[1] + [200] + [1], 
                 DNN_size=[1] + [20] * 3 + [1]):
        if improved:
            k = 65.0
            NumTrain=380
        self.a = a
        self.b = b
        self.k = k
        super().__init__(name="Helmholtz", v=v, NumTrain=NumTrain, NumTest=NumTest, max_iter=max_iter, FP_size=FP_size, DNN_size= DNN_size)

    def func(self, x):
        #print(f"x:{x},result:{np.sin(self.k * x)}")
        return np.sin(self.k * x)
        
    def gen_pde(self):
        def pde(x, y):
            dy_xx = dde.grad.hessian(y, x, i=0, j=0)
            return dy_xx + self.k ** 2 * y
        return pde

    def gen_geomtime(self):
        return dde.geometry.Interval(self.a, self.b)
    
    def gen_data(self):
        bc = dde.icbc.DirichletBC(self.geomtime, self.func, lambda _,on_boundary:on_boundary)
        data = dde.data.PDE(self.geomtime, self.pde, bc, num_domain=self.NumTrain, num_boundary=2, num_test=self.NumTest, solution=self.func, train_distribution='uniform')
        return data
    
    def plot(self, model=None, title='Solution'):
        x = np.linspace(self.a, self.b, 2000)
        x = np.expand_dims(x, axis=1)
        y = self.func(x)
        plt.plot(x, y, label="Ground truth", alpha=0.5)

        if model is not None:
            y_pred = model.model.predict(x)
            plt.plot(x, y_pred, label="Prediction", alpha=0.5)
            plt.plot(x, y_pred - y, label="Error", alpha=0.5)
        
        plt.title(title)
        plt.xlabel("x")
        plt.ylabel("y")
        plt.legend()
        plt.show()

class HighFrequencyCase2D1(cFPCT_DNN_Cases):
    def __init__(self,
                 v=4.6, 
                 NumTrain=10000, 
                 NumTest=20000, 
                 max_iter=40000, 
                 FP_size=[2] + [200] + [1], 
                 DNN_size=[2] + [20] * 3 + [1]):
        super().__init__(name="HighFrequencyCase3", v=v, NumTrain=NumTrain, NumTest=NumTest, max_iter=max_iter, FP_size=FP_size, DNN_size= DNN_size)

    def func(self, x):
        return np.cos(50*x[:,0] + 50*x[:,1])

    def gen_geomtime(self):
        return dde.geometry.Rectangle([-1, -1], [1, 1])
    
    def plot_heatmap_at_axes(self, X, y, axes, title):
        axes.set_title(title)
        axes.set_xlabel('x')
        axes.set_ylabel('y')
        axes.pcolormesh(X[0,:,:], X[1,:,:], y, cmap='rainbow')
    
    def plot(self, model=None, title='Solution'):
        num = 1000
        x=np.linspace(-1.0, 1.0, num)
        y=np.linspace(-1.0, 1.0, num)
        X,Y =np.meshgrid(x,y)
        X = np.stack([X.reshape(-1), Y.reshape(-1)], axis=1)
        y = self.func(X)[:,np.newaxis]
        if model is None:
            X = np.stack([X[:, 1].reshape(num, num), X[:, 0].reshape(num, num)], axis=0)
            fig, axes = plt.subplots(1, 1)
            self.plot_heatmap_at_axes(X, y.reshape(num, num), axes=axes, title='Exact solution')
        else:
            fig, axes = plt.subplots(1, 3, figsize=(15, 5))
            model_y = model.model.predict(X)
            X = np.stack([X[:, 1].reshape(num, num), X[:, 0].reshape(num, num)], axis=0)
            self.plot_heatmap_at_axes(X, y.reshape(num, num), axes=axes[0], title='Exact solution')
            self.plot_heatmap_at_axes(X, model_y.reshape(num, num), axes[1], title=title)
            self.plot_heatmap_at_axes(X, np.abs(model_y - y).reshape(num, num) , axes[2], title='Absolute error')
        plt.show()
        return fig, axes